# %%
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import os

# %%
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

batch_size = 4

trainset = torchvision.datasets.CIFAR100(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)


testset = torchvision.datasets.CIFAR100(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                         shuffle=False, num_workers=2)

classes = ('beaver', 'dolphin', 'otter', 'seal', 'whale',
           'aquarium fish', 'flatfish', 'ray', 'shark', 'trout',
           'orchids', 'poppies', 'roses', 'sunflowers', 'tulips',
           'bottles', 'bowls', 'cans', 'cups', 'plates',
           'apples', 'mushrooms', 'oranges', 'pears', 'sweet',
           'peppers',
           'clock', 'computer keyboard', 'lamp', 'telephone',
           'television',
           'bed', 'chair', 'couch', 'table', 'wardrobe',
           'bee', 'beetle', 'butterfly', 'caterpillar', 'cockroach',
           'bear', 'leopard', 'lion', 'tiger', 'wolf',
           'bridge', 'castle', 'house', 'road', 'skyscraper',
           'cloud', 'forest', 'mountain', 'plain', 'sea',
           'camel', 'cattle', 'chimpanzee', 'elephant', 'kangaroo',
           'fox', 'porcupine', 'possum', 'raccoon', 'skunk',
           'crab', 'lobster', 'snail', 'spider', 'worm',
           'baby', 'boy', 'girl', 'man', 'woman',
           'crocodile', 'dinosaur', 'lizard', 'snake', 'turtle',
           'hamster', 'mouse', 'rabbit', 'shrew', 'squirrel',
           'maple', 'oak', 'palm', 'pine', 'willow',
           'bicycle', 'bus', 'motorcycle', 'pickup truck', 'train',
           'lawn-mower', 'rocket', 'streetcar', 'tank', 'tractor')

# %%

import matplotlib.pyplot as plt
import numpy as np

# functions to show an image


def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


# get some random training images
dataiter = iter(trainloader)
images, labels = dataiter.next()
print(images[0].shape)

# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))


# %%
class channel_attention(nn.Module):
    def __init__(self, channel, ratio=4):
        super(channel_attention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.shared_mlp = nn.Sequential(
            nn.Conv2d(channel, channel // ratio, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(channel // ratio, channel, 1, bias=False)
        )
        self.sigmond = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.shared_mlp(self.avg_pool(x))
        max_out = self.shared_mlp(self.max_pool(x))
        return self.sigmond(avg_out + max_out)

# %%
class spatial_attention(nn.Module):
    def __init__(self):
        super(spatial_attention, self).__init__()
        self.conv = nn.Conv2d(2, 1, 7, padding=3)
        self.sigmond = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        out = torch.cat([avg_out, max_out], dim=1)
        out = self.conv(out)
        out = self.sigmond(out)
        return out

# %%
class cbam_block(nn.Module):
    def __init__(self, channel):
        super(cbam_block, self).__init__()
        self.c_a = channel_attention(channel)
        self.s_a = spatial_attention()
    
    def forward(self, x):
        out = self.c_a(x) * x
        out = self.s_a(out) * out
        return out

# %%
class res_block(nn.Module):
    def __init__(self, in_c, out_c, stride):
        self.rec_stride = stride
        super(res_block, self).__init__()
        self.left = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=3, stride=stride, padding=1, groups=in_c, bias=False),
                                    nn.BatchNorm2d(out_c),
                                    nn.ReLU(),
                                    nn.Conv2d(out_c, out_c, kernel_size=1, stride=1, padding=0, bias=False),
                                    nn.BatchNorm2d(out_c)
                                    )
        self.shortcut = nn.Sequential()
        if stride != 1 or in_c != out_c:
            self.shortcut = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=1, stride=stride, bias=False),
                                            nn.BatchNorm2d(out_c)
                                            )
        self.cbam = cbam_block(out_c)
        
    def forward(self, x):
        #print('x:', x.shape)
        out = self.left(x)
        out = self.cbam(out)
        #print('out:', out.shape, self.rec_stride)
        #print('shortcut:', self.shortcut(x).shape)
        out += self.shortcut(x)
        out = F.relu(out)
        return out


# %%
class Net(nn.Module):
    def __init__(self, res_b, num_class = 100):
        super().__init__()
        self.inchannel = 32
        self.conv1 = nn.Sequential(
                            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1, bias=False),
                            nn.BatchNorm2d(32),
                            nn.ReLU(),
                            )
        self.layer1 = self.make_layer(res_b, 64, 2, stride=2)
        #self.layer2 = self.make_layer(res_b, 32, 2, stride=2)
        #self.layer3 = self.make_layer(res_b, 128, 2, stride=2)
        self.layer4 = self.make_layer(res_b, 128, 2, stride=2)
        self.fc = nn.Linear(128 * 4 * 4, num_class)

    def make_layer(self, block, channels, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for sss in strides:
            layers.append(block(self.inchannel, channels, sss))
            self.inchannel = channels 
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv1(x)
        out = self.layer1(out)
        #out = self.layer2(out)
        #out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 2)
        #print(out.shape)
        #out = out.view(out.size(0), -1)
        out = torch.flatten(out, 1)
        out = self.fc(out)
        return out

net = Net(res_block)

# %%
import torch.optim as optim
from tqdm import tqdm

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)

losses = []
acc = []

# %%
for epoch in range(4):  # loop over the dataset multiple times

    running_loss = 0.0
    cur_acc = 0
    for i, data in tqdm(enumerate(trainloader, 0)):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        mytensor = torch.tensor(inputs).to(device)
        # forward + backward + optimize
        outputs = net(mytensor)
        outputs = outputs.to(torch.device("cpu"))
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        _, pred = torch.max(outputs, 1)
        for j in range(len(labels)):
            if pred[j] == labels[j]:
                cur_acc += 1

        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            losses.append(running_loss / 2000)
            acc.append(cur_acc / 2000 / 4)
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0
            cur_acc = 0

print('Finished Training')

# %%
PATH = './cifar100_cbam_resnet.pth'
torch.save(net.state_dict(), PATH)


# %%
plt.figure()
plt.title("loss")
plt.plot(losses)
plt.figure()
plt.title("accuracy")
plt.plot(acc)

# %%
dataiter = iter(testloader)
images, labels = dataiter.next()

# print images
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))

# %%
net = Net(res_block)
net.load_state_dict(torch.load(PATH))
_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}'
                              for j in range(4)))


# %%
correct = 0
total = 0
# since we're not training, we don't need to calculate the gradients for our outputs
with torch.no_grad():
    for data in testloader:
        images, labels = data
        # calculate outputs by running images through the network
        outputs = net(images)
        # the class with the highest energy is what we choose as prediction
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')

# %%
# prepare to count predictions for each class
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}

# again no gradients needed
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predictions = torch.max(outputs, 1)
        # collect the correct predictions for each class
        for label, prediction in zip(labels, predictions):
            if label == prediction:
                correct_pred[classes[label]] += 1
            total_pred[classes[label]] += 1


# print accuracy for each class
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / (total_pred[classname] + (1 if total_pred[classname] == 0 else 0))
    print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')


# %%
l_index = [i for i in range(len(losses))]
plt.plot(l_index, losses)


